import pandas as pd
from numpy import *
import time

def prepare_data():
    df = pd.read_excel("./超市数据集.xls")
    prd_id = array(df.columns)
    prd_name = array(df.ix[0])
    prd_dict = dict(zip(prd_id, prd_name))

    data = df.iloc[1:,:].as_matrix()
    items_arr = []

    for i in range(0, len(data)):
        row = data[i]
        new_row = []
        for j in range(0, len(row)):
            if(row[j] == 'T'):
                new_row.append(prd_name[j])
        items_arr.append(new_row)

    new_ids = []
    for id in prd_name:
        if not [id] in new_ids:
            new_ids.append([id])
    
    prd_items = map(frozenset, new_ids)
    item_list = list(prd_items)

    data_set = list(map(set, items_arr))
    return item_list, data_set

def scan(data_set,item_list, min_sup):
    cnt = {}
    for row in data_set:
        for item in item_list:
            if item.issubset(row):
                if not item in cnt:
                    cnt[item] = 1
                else:
                    cnt[item] += 1
    d_len = float(len(data_set))
    ret_list = []
    support_data = {}
    for key in cnt:
        support = cnt[key]/d_len
        if(support > min_sup):
            ret_list.insert(0, key)
        support_data[key] = support
    return ret_list, support_data

def apriori_gen(in_list, k):
    ret_list = []
    list_len = len(in_list)
    for i in range(list_len):
        for j in range(i+1, list_len):
            l1 = list(in_list[i])[:k-2]
            l2 = list(in_list[j])[:k-2]
            l1.sort()
            l2.sort()
            if l1 == l2:
                ret_list.append(in_list[i]|in_list[j])
    return ret_list

def apriori(data_set, item_list, min_sup = 0.4):
    l1, support_data = scan(data_set, item_list, min_sup)
    L = [l1]
    k = 2

    while len(L[k-2]) > 0:
        ap_gen = apriori_gen(L[k-2], k)
        l_k, sup_k = scan(data_set, ap_gen, min_sup)
        support_data.update(sup_k)
        L.append(l_k)
        k += 1
    return L, support_data

def calc_conf(freq_set, H, support_data, brl, min_conf = 0.7):
    prunedH = []
    for conseq in H:
        conf = support_data[freq_set]/support_data[freq_set - conseq]
        if(conf >= min_conf):
            #print(freq_set-conseq, "-->", conseq, "conf:", conf)
            brl.append((freq_set-conseq, conseq, conf))
            prunedH.append(conseq)
    return prunedH

def rules_from_conseq(freq_set, H, support_data, brl, min_conf=0.7):
    m = len(H[0])
    if (len(freq_set) > (m + 1)):
        Hmp1 = apriori_gen(H, m + 1)
        Hmp1 = calc_conf(freq_set, Hmp1, support_data, brl, min_conf)
        if (len(Hmp1) > 1):
            rules_from_conseq(freq_set, Hmp1, support_data, brl, min_conf)

def gen_rules(l, support_data, min_conf = 0.7):
    rule_list = []
    for i in range(1, len(l)):
        for freq_set in l[i]:
            H1 = [frozenset([item]) for item in freq_set]
            if (i>1):
                rules_from_conseq(freq_set, H1, support_data, rule_list, min_conf)
            else:
                calc_conf(freq_set, H1, support_data, rule_list, min_conf)
    return rule_list


item_list, data_set = prepare_data()

begin_time = time.time()
l, support_data = apriori(data_set, item_list, 0.3)
rules = gen_rules(l, support_data, min_conf = 0.5)

end_time = time.time()

print("-------------------------Rules-----------------------")
for rule in rules:
    print(rule)
print("-------------------------Rules-----------------------")
print("Elapsed Time:" + str((end_time - begin_time)*1000) + " milliseconds")
